#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File  :  1_ali_to_coco.py
@Time  :  2022/05/20 19:57:23
@Author :  Xu Qian
@Version :  1.0.0
@Contact :  Qian.xu3@cn.bosch.com
@Desc  :  None
'''
import os
from pycocotools.coco import COCO
import json
import mmcv
from pathlib import Path
import csv
import cv2
import numpy as np
from tqdm import tqdm
import math

import sys

sys.path.append(os.path.dirname(__file__))

categories_box = [
    {'supercategory': 'slot_id', 'id': 1, 'name': 'slot_id'},
    {'supercategory': 'others', 'id': 2, 'name': 'others'},
]


def get_category_id(label_name):
    for cat in categories_box:
        if cat['name'] == label_name:
            return cat['id']


class CocoConverter:

    def __init__(self,
                 csv_file,
                 # json_file,
                 # old_json,
                 result_path,
                 txt_path,
                 ) -> None:
        # self.json_file = json_file
        # self.raw_info_json = self.load_json(self.json_file)
        # self.old_json = self.load_single_json(old_json)
        self.txt_path = txt_path
        self.csv_file = csv_file
        self.raw_info = self.load_csv(csv_file)
        self.width = 1000
        self.height = 1000
        self.result_path = result_path

        self.gen_coco_box()

    def load_img_path(self, imgs_path):
        imgs_path_glob = Path(imgs_path).rglob("*.png")
        img_name_path_dict = {x.name: str(x) for x in imgs_path_glob}

        return img_name_path_dict

    def load_img(self, img_name_path_dict, file_name):
        img_path = img_name_path_dict[file_name]
        img = cv2.imread(img_path)
        return img

    def load_list(self, list_path):
        test_img_list = list()
        with open(list_path, 'r') as f:
            img_list = f.readlines()
        for img_path in img_list:
            img_path = img_path.strip()
            test_img_list.append(img_path)
        return test_img_list

    def load_json(self, label_dir):
        result_sum = list()
        labels_path_glob = Path(label_dir).rglob("*.json")
        labels_path_list = list(labels_path_glob)
        for raw_info in tqdm(labels_path_list):
            with open(str(raw_info), 'r') as f:
                result = json.load(f)
                result_sum.append(result)
        return result_sum

    def load_single_json(self, json_path):
        with open(str(json_path), 'r') as f:
            result = json.load(f)
        return result

    def merge_raw_info(self, json_result, csv_result):
        json_result.extend(csv_result)
        return json_result

    def get_coco_box(self, box):
        assert len(box) == 4
        min_x = min_y = 100000
        max_x = max_y = 0
        min_x = min(min_x, min(box[0:4:2]))
        min_y = min(min_y, min(box[1:4:2]))
        max_x = max(max_x, max(box[0:4:2]))
        max_y = max(max_y, max(box[1:4:2]))
        return [min_x, min_y, max_x - min_x, max_y - min_y]

    def get_target_box(self, box):
        bboxps = np.array(box, dtype=np.float32).reshape((-1, 2))
        rbbox = cv2.minAreaRect(bboxps)
        x, y, w, h, a = rbbox[0][0], rbbox[0][1], rbbox[1][0], rbbox[1][1], rbbox[2]
        if w < 2 or h < 2:
            return
        while not 0 < a <= 90:
            if a == -90:
                a += 180
            else:
                a += 90
                w, h = h, w
        a = a / 180 * np.pi
        assert 0 < a <= np.pi / 2
        return x, y, w, h, a

    def load_csv(self, label_file):
        raw_info = csv.DictReader(open(label_file, mode='r'))
        raw_info = list(raw_info)
        return raw_info

    def parse_images(self, start_id):
        # cur_id = 1
        cur_id = 0
        name_id_dict = {}
        img_info_list = []
        # old_img_raw = self.old_json['images']
        # start_id = old_img_raw[-1]['id']
        for single_raw_info in self.raw_info:
            file_name = single_raw_info['url']
            file_name_list = str(file_name).split("/")
            for single_part in file_name_list:
                if '8328C' in single_part or '7185C' in single_part:
                    file_name = os.path.join("/home/ych6szh/Documents/parking_structure/ipm_marker_new/slot_img",
                                             single_part, file_name_list[-1])
            if '.png' not in file_name:
                continue
            print(file_name)
            if file_name in name_id_dict.keys():
                continue
            id = start_id + cur_id
            cur_id += 1
            # img = self.load_img(self.img_name_path_dict, file_name)
            # height,width,_ = img.shape
            self.height = 1000
            self.width = 1000

            img_info = {
                'license': 0,
                'file_name': file_name,
                'height': self.height,
                'width': self.width,
                'id': id,
            }
            img_info_list.append(img_info)

            name_id_dict[file_name] = id

        return img_info_list, name_id_dict

    def parse_annotations_box(self, raw_info, name_id_dict, start_id):
        annotations = []
        cur_id = 0
        old_ann_list = []

        file = open(self.txt_path, 'w')

        for single_raw_info in tqdm(raw_info, desc='parsing lines ...'):
            try:
                file_name = Path(single_raw_info['url'])
                name = Path(single_raw_info['url']).name
                file_name_list = str(file_name).split("/")
                for single_part in file_name_list:
                    if '8328C' in single_part or '7185C' in single_part:
                        file_name_new = os.path.join("/home/ych6szh/Documents/parking_structure/ipm_marker_new/slot_img",
                                                 single_part, file_name_list[-1])
                if '.png' not in file_name_new:
                    continue
            except:
                continue
            try:
                anno_json = json.loads(single_raw_info['sd_result'])
            except:
                print("skip img is {}".format(name))
                file.write(name + '\n')
                continue

            # id of the image
            if name_id_dict.get(file_name_new) is None:
                print(file_name_new)
                if isinstance(anno_json, int):
                    continue
                if isinstance(anno_json, str):
                    continue
                for anno_item in anno_json['items']:
                    if 'box' in anno_item['labels'].keys():
                        # print(anno_item['labels']['line'])
                        if anno_item['labels']['box'] == 'slot_id':
                            print("error!")
                continue
            else:
                if file_name_new in old_ann_list:
                    continue
                image_id = name_id_dict[file_name_new]
                old_ann_list.append(file_name_new)
            if isinstance(anno_json, int):
                continue
            if isinstance(anno_json, str):
                continue
            for anno_item in anno_json['items']:
                anno_type = anno_item['meta']['type']
                if anno_type == 'PARAOG':
                    bbox = anno_item['meta']['geometry']

                    bboxps = np.array(bbox, dtype=np.float32).reshape((-1, 2))
                    rbbox = cv2.minAreaRect(bboxps)
                    box = cv2.boxPoints(rbbox)
                    box = np.int0(box)
                    # if anno_item['labels']['box'] == 'ignore':
                    #     cv2.drawContours(img, [box], 0, (0, 0, 255), 2)
                    # else:
                    #     cv2.drawContours(img, [box], 0, (0, 255, 0), 2)

                    bbox = self.get_target_box(bbox)
                    if bbox is None:
                        continue
                    area = bbox[2] * bbox[3]
                    # id of the box
                    id = start_id + cur_id
                    cur_id += 1
                    iscrowd = 0
                    # if the box is ignored, set the categories as others and \
                    # iscrowd as 1
                    if anno_item['labels']['box'] == 'ignore':
                        bbox_type = 'others'
                        iscrowd = 1
                    else:
                        bbox_type = anno_item['labels']['box']
                        c_tag = anno_item['labels']['c-tag']
                        d_tag = anno_item['labels']['d-tag']
                        assert c_tag in ['clear', 'blurry', 'severe_blurry']
                        assert d_tag in ['intact', 'occlude', 'severe_occlude']
                        if c_tag == 'severe_blurry' or \
                                d_tag == 'severe_occlude':
                            iscrowd = 1
                    category_id = get_category_id(bbox_type)

                    if category_id is None:
                        print(bbox_type)
                        continue

                    anno = {
                        'bbox': bbox,
                        'area': area,
                        'iscrowd': iscrowd,
                        'image_id': image_id,
                        'category_id': category_id,
                        'id': id,
                    }
                    annotations.append(anno)

            # cv2.imwrite(save_dir, img)
        file.close()
        return annotations

    def gen_coco_box(self, ):

        images, name_id_dict = self.parse_images(start_id=0)
        annotations = self.parse_annotations_box(self.raw_info, name_id_dict, start_id=0)
        struct_json_boxes = {
            'images': images,
            'annotations': annotations,
            'categories': categories_box,
        }

        result_file_path = Path(self.result_path).joinpath('box_train_new_all.json')
        json.dump(struct_json_boxes, open(str(result_file_path), "w"), indent=4)


def main():
    json_file = r'/home/ych6szh/Documents/parking_structure/slot_train_0212/slot_ali'
    old_json = r'/home/ych6szh/Documents/parking_structure/slot_train_0212/box_train_new_0212_baidu.json'
    csv_file = r'/home/ych6szh/Documents/parking_structure/ipm_marker_new/IPM_ROAD_MARKER_0226.csv'
    imgs_path = r'/mnt/sdb/share/15_offline_raw_data/'
    result_path = r'/home/ych6szh/Documents/parking_structure/slot_train_0212/'
    txt_path = r'/home/ych6szh/Documents/parking_structure/skip_img_list.txt'
    img_save_path = r'/home/ych6szh/Documents/parking_structure/ipm_marker_new/train_slot_id_stop_line/label_show/'
    ccc = CocoConverter(csv_file, result_path, txt_path)


if __name__ == "__main__":
    main()
